Context Managers

Context managers allow you to call setup code before a block of code is executed and teardown when the code is done. They are very useful in resource management but handy in many other situations.

A context manager object has a an __enter__(self) method that is called at start of the with block and an __exit__(self, exc_type, exc_value, traceback) that is called at the end of the with block. (The extra arguments to __exit__ are in case of exception).

with open('/path/to/file.txt') as fo:
    process(fo)

is very much like the following ("double humped code" as Raymond calls it :)

fo = open('/path/to/file.txt')
try:
    process(fo)
finally:
    fo.close()

API

A context manager object should have the following methods.

enter(self)

Called before the code block inside the with is executed. Can return an object which is bound in the as clause.

exit(self, exc_type, exc_value, traceback)

Called after the code block insided the with is executed (even when an exception is raised). Last 3 parameters are in case of an error, on normal execution they will be None.

If __exit__ returns True, exceptions are "swallowed".

Nesting

Context managers can be nested (separated with ,). The __enter__ will be called in order, the __exit__ in reverse order.



In [1]:

    
class echo(object):
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print('enter {}'.format(self.name))

    def __exit__(self, exc_type, exc_value, traceback):
        print('exit {}'.format(self.name))


with echo('a'), echo('b'), echo('c'):
    pass









    



enter a
enter b
enter c
exit c
exit b
exit a

Notes

Context managers are in use in the standard library (file objects, locks ...), and outside.
contextlib in the standard library help with writing context managers (more on that later).

Exercise: Greeter

Write a context manager that print 'Hai' at the start of the block and 'Bai' at end.



In [6]:

    
class greeter(object):
    def __enter__(self):
        print('Hai')

    def __exit__(self, exc_type, exc_value, traceback):
        print('Bai')



In [7]:

    
with greeter():
    print('Wassup?')
# Should print
# Hai
# Wassup?
# Bai









    



Hai
Wassup?
Bai

Exercise: Timing

Write a context manager that times the code inside the with statement, it should get a name argument to print.



In [8]:

    
from time import time


class timed_block(object):
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        self.start = time()

    def __exit__(self, exc_type, exc_value, traceback):
        duration = time() - self.start
        print('{} took {:0.2f}sec'.format(self.name, duration))



In [9]:

    
from time import sleep
with timed_block('sleep'):
    sleep(0.2)
# Should print 'sleep took 0.2sec'









    



sleep took 0.20sec

Exercise: Closer

Write a context manager that closes the object it is passed by calling obj.close().



In [11]:

    
class closing(object):
    def __init__(self, obj):
        self.obj = obj

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_value, traceback):
        self.obj.close()



In [13]:

    
class Stream(object):
    def close(self):
        print('Closing stream')
        
stream = Stream()
with closing(stream):
    pass
# Should print 'Closing stream'









    



Closing stream

Exercise: Databases

Write a context manager that gets a database connection, provides a cursor to the with block and then either commits if everything went well, otherwise rollbacks.

See DB2 API for database API in Python.

However, all you need to know is that a connection cursor method will create a new cursor. e.g.:

import sqlite3
conn = sqlite3.connect(':memory:')
cur = conn.cursor()



In [14]:

    
class dbctx(object):
    def __init__(self, conn):
        self.conn = conn

    def __enter__(self):
        return self.conn.cursor()

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            print('Committing')
            self.conn.commit()
        else:
            print('Rolling back')
            self.conn.rollback()



In [16]:

    
import sqlite3
conn = sqlite3.connect(':memory:')
with dbctx(conn) as cur:
    cur.execute('SELECT 1')
# commit

with dbctx(conn) as cur:
    cur.execute('SELECT * FROM whatever')
# rollback









    



---------------------------------------------------------------------------
OperationalError                          Traceback (most recent call last)
<ipython-input-16-ee366a6bdc16> in <module>()
      6 
      7 with dbctx(conn) as cur:
----> 8     cur.execute('SELECT * FROM whatever')
      9 # rollback

OperationalError: no such table: whatever





    



Committing
Rolling back

contextlib

contextlib provides some functions that help with writing context managers. Notably the contextmanager decorator.

Exercise: contextlib

Rewrite exercise 3 using the contextlib decorator.



In [ ]:

    
from contextlib import contextmanager


@contextmanager
def dbctx(conn):
    try:
        yield conn.cursor()
        conn.commit()
    except:
        conn.rollback()
        raise  # Need to re-raise